import sys
import os
# 获取父目录的绝对路径
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
# 将父目录添加到sys.path
sys.path.append(parent_dir)
from utils.log import logger
from utils.chatglm_utils import ChatGLMUtils
from modules.agg_recognition import AggRecognition
if __name__ == "__main__":
    glm_util = ChatGLMUtils()
    # 使用chatglm2只需要修改model_path和checkpoint（改到glm2的路径），cuda_index就是cuda的序号。
    # glm_util.load_model(local_loading=True,
    #                     model_path=f"/home/cike/smart/GLM_agg/outputV4/agg-chatglm-6b-lora.pt")
    glm_util.load_model(local_loading=True, 
                    model_path="/home/cike/ytc/GLM2/checkpoints/Lora7_25_combine/checkpoint-64000/pytorch_model.bin",
                    cuda_index=2)
    agg_reg=AggRecognition(glm_util,"/home/cike/ytc/val-reg-with-glm/data/data_dict.json")
    # glm_util.load_model(local_loading=True,
    #                     # model_path=f"/home/cike/smart/val-reg-with-glm/models/all-chatglm-6b-lora.pt")
    #                     model_path=f"/home/cike/smart/val-reg-with-glm/models/pytorch_model.bin")
    # agg_reg=AggRecognition(glm_util,"/home/cike/smart/val-reg-with-glm/data/data_dict.json")
    test={
    "query":"去年12月合同金额总和大于500万的分部有多少？",
    "sel":[{"id" : "AUGMENTED_DATASET_LEVEL.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.3943e92bed35dac104c3a1f2b2195a5e-LEVEL-1667896190875","name" : "年","type" : "STRING","timeLevel" : "year","member_amount" : 1},{"id" : "AUGMENTED_DATASET_LEVEL.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_DIMENSION_FIELD.affebb66d15867660b4dcd4712c9a7d9-LEVEL-1667896197688","name" : "月","type" : "STRING","timeLevel" : "month","member_amount" : 0},{"id" : "AUGMENTED_DATASET_MEASURE.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_METRICS.I8a8ae5ca0178549554951b9501785d983aaa005e","name" : "合同金额","type" : "DOUBLE","timeLevel" : "","member_amount" : 0}],
    "measure": ["AUGMENTED_DATASET_MEASURE.I8a8ae5ca0178549554951b9501785cefe3f00058.MT_METRICS.I8a8ae5ca0178549554951b9501785d983aaa005e"]
    }
    print(agg_reg.recognition(test['query'],test['sel'],test['measure']))




